import torch
import numpy as np
from torchvision import datasets, transforms
from matplotlib import pyplot as plt
from DenoriseDiffusion import DoubleUnet, Diffusion
from tqdm import tqdm
def im_show(arr, row, col):
    fig, axs = plt.subplots(nrows=row, ncols=col)
    for r in range(row):
        for c in range(col):
            axs[r][c].imshow(arr[r*col+c])
            axs[r][c].axis('off')
    plt.show()
transform = transforms.ToTensor()
train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=len(train_set), shuffle=False)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=len(test_set), shuffle=False)
for images, labels in train_loader:
    train_images, train_labels = images, labels
    break
for images, labels in test_loader:
    test_images, test_labels = images, labels
    break
x = torch.cat((train_images, test_images), dim=0)
# 以下为模板
device = 'cuda' if torch.cuda.is_available() else 'cpu'
in_channels = 1
image_shape=(28,28)
beta=(1e-4,0.02)
step=400
feat=16
unet = DoubleUnet(in_channels=in_channels, image_shape=image_shape, step=step, feat=feat)
diffusion = Diffusion(model=unet, beta=beta, image_shape=(in_channels, *image_shape), step=step, device=device).to(device)
opt = torch.optim.Adam(params=diffusion.model.parameters(), lr=1e-4)
batch_size=256
length=len(x)
Epoch = 400
for epoch in range(Epoch):
    for i in tqdm(range(0, length, batch_size), desc=f"Epoch: {epoch}"):
        tx = x[i:i+batch_size].to(device)
        diffusion.train()
        opt.zero_grad()
        loss = diffusion(tx)
        loss.backward()
        opt.step()
    if epoch % 20 == 0:
        diffusion.eval()
        images = diffusion.sample(16).cpu().detach().numpy().transpose(0,2,3,1)
        im_show(images, 4, 4)
torch.save(diffusion, "./model/model.pth")